/**
 * 
 */
package edu.ubc.javascript;

import static edu.ubc.javascript.ASTFactory.name;

import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import com.google.javascript.jscomp.mozilla.rhino.Node;
import com.google.javascript.jscomp.mozilla.rhino.Token;
import com.google.javascript.jscomp.mozilla.rhino.ast.ArrayLiteral;
import com.google.javascript.jscomp.mozilla.rhino.ast.Assignment;
import com.google.javascript.jscomp.mozilla.rhino.ast.AstNode;
import com.google.javascript.jscomp.mozilla.rhino.ast.Block;
import com.google.javascript.jscomp.mozilla.rhino.ast.ExpressionStatement;
import com.google.javascript.jscomp.mozilla.rhino.ast.FunctionCall;
import com.google.javascript.jscomp.mozilla.rhino.ast.FunctionNode;
import com.google.javascript.jscomp.mozilla.rhino.ast.IfStatement;
import com.google.javascript.jscomp.mozilla.rhino.ast.InfixExpression;
import com.google.javascript.jscomp.mozilla.rhino.ast.Name;
import com.google.javascript.jscomp.mozilla.rhino.ast.NodeVisitor;
import com.google.javascript.jscomp.mozilla.rhino.ast.ObjectLiteral;
import com.google.javascript.jscomp.mozilla.rhino.ast.ObjectProperty;
import com.google.javascript.jscomp.mozilla.rhino.ast.PropertyGet;
import com.google.javascript.jscomp.mozilla.rhino.ast.ReturnStatement;
import com.google.javascript.jscomp.mozilla.rhino.ast.UnaryExpression;
import com.google.javascript.jscomp.mozilla.rhino.ast.VariableDeclaration;
import com.google.javascript.jscomp.mozilla.rhino.ast.VariableInitializer;

public class TACVisitor implements NodeVisitor {

	private ReflectiveAstTransformer tx;
	
	public TACVisitor(ReflectiveAstTransformer tx) {
		this.tx = tx;
	}
	
	private int count = 0;
	
	private String newName() {
		return "_r" + (count++);
	}
	
	private Set<Node> done = new HashSet<Node>();
	
	private Node transform(boolean lvalue, AstNode expr, List<Node> acc, List<Node> gs) {
		if(expr == null) {
			return expr;
		}
		if(ContextCheck.isAssignmentOp(expr.getType())) {
			InfixExpression infix = (InfixExpression) expr;
			Node left = transform(true, infix.getLeft(), acc, gs); 
			Node right = transform(false, infix.getRight(), acc, gs);
			infix.setRight((AstNode) right);
			infix.setLeft((AstNode) left);
			Name name = ASTFactory.name(newName());
			Assignment assign = ASTFactory.assign(name("var " + name.getString()), infix);
			acc.add(ASTFactory.exprStmt(assign));
			return name;
		} else if(expr instanceof InfixExpression){
			InfixExpression infix = (InfixExpression) expr;
			if(infix.getOperator() == Token.AND) {
				
				Node left = transform(lvalue, infix.getLeft(), acc, gs); 
				
				List<Node> rightAcc = new LinkedList<Node>();
				Node right = transform(lvalue, infix.getRight(), rightAcc, gs);
				Name name = ASTFactory.name(newName());
				Block b = new Block();
				for(Node n : rightAcc) {
					b.addStatement((AstNode) n);
				}
				b.addStatement(ASTFactory.exprStmt(ASTFactory.assign(name("var " + name.getString()), (AstNode) right)));
				IfStatement ifStatement = ASTFactory.ifThenElse((AstNode) left, b, ASTFactory.exprStmt(ASTFactory.assign(name("var " + name.getString()), ASTFactory.name("false"))));
				
				acc.add(ifStatement);
				return name;
			} else if(infix.getOperator() == Token.OR) {
				
				Node left = transform(lvalue, infix.getLeft(), acc, gs); 
				
				List<Node> rightAcc = new LinkedList<Node>();
				Node right = transform(lvalue, infix.getRight(), rightAcc, gs);
				Name name = ASTFactory.name(newName());
				Block b = new Block();
				for(Node n : rightAcc) {
					b.addStatement((AstNode) n);
				}
				b.addStatement(ASTFactory.exprStmt(ASTFactory.assign(name("var " + name.getString()), (AstNode) right)));
				IfStatement ifStatement = ASTFactory.ifThenElse((AstNode) left, ASTFactory.exprStmt(ASTFactory.assign(name("var " + name.getString()), (AstNode) left)), b);
	
				acc.add(ifStatement);
				return name;
			}
			else if(lvalue && ContextCheck.isAtomicExpression(infix.getRight())) {
				Node left = transform(false, infix.getLeft(), acc, gs); 
				infix.setLeft((AstNode) left);
				if(!done.contains(infix) && (infix.getType() == Token.GETPROP)) {
					AstNode nameNode = infix.getRight();
					if(nameNode instanceof Name) {
						ExpressionStatement exprStmt = ASTFactory.exprStmt(ASTFactory.call("___setter", (AstNode) left, ASTFactory.string(nameNode.getString())));
						gs.add(exprStmt);
						done.add(infix);
					}
					//if(!obj._old) {
					//  old._old = clone(obj);
					//  _tx.writers.push(obj);
					//}
					//AstNode exprStmt = ASTFactory.exprStmt(ASTFactory.assign(ASTFactory.get((AstNode) left, "_dirty"), ASTFactory.name("true")));
					//IfStatement ifStatement = ASTFactory.ifThen(ASTFactory.get((AstNode) left, "_id===undefined"), exprStmt);
				}
				return infix;
			} else {
				Node left = transform(lvalue, infix.getLeft(), acc, gs); 
				Node right = transform(lvalue, infix.getRight(), acc, gs);
				infix.setRight((AstNode) right);
				infix.setLeft((AstNode) left);
				Name name = ASTFactory.name(newName());
				Assignment assign = ASTFactory.assign(name("var " + name.getString()), infix);
				if(infix.getType() == Token.GETPROP) {
					//AstNode exprStmt = ASTFactory.exprStmt(ASTFactory.call("_stack.push", ASTFactory.name("'"+name.getString() +"'")));
					if(infix.getRight() instanceof Name) {
						ExpressionStatement exprStmt = ASTFactory.exprStmt(ASTFactory.call("___getter", name, ASTFactory.string(infix.getRight().getString())));
						gs.add(exprStmt);
					}
				}
				acc.add(ASTFactory.exprStmt(assign));
				return name;
			}
		} else if(expr instanceof PropertyGet) {
			PropertyGet pg = (PropertyGet) expr;
			Node left = transform(lvalue, pg.getLeft(), acc, gs); 
			Node right = transform(lvalue, pg.getRight(), acc, gs);
			pg.setRight((AstNode) right);
			pg.setLeft((AstNode) left);
			Name name = ASTFactory.name(newName());
			Assignment assign = ASTFactory.assign(name("var " + name.getString()), pg);
		
			acc.add(ASTFactory.exprStmt(assign));
			return name;
		} /*else if(expr instanceof ElementGet) {
			ElementGet eg = (ElementGet) expr;
			Node left = transform(lvalue, eg.getTarget(), acc); 
			Node right = transform(lvalue, eg.getElement(), acc);
			eg.setElement((AstNode) right);
			eg.setTarget((AstNode) left);
			Name name = ASTFactory.name(newName());
			Assignment assign = ASTFactory.assign(name("var " + name.getString()), eg);
			acc.add(ASTFactory.exprStmt(assign));
			return name;
		}*/
		else if(expr instanceof FunctionCall) {
			FunctionCall call = (FunctionCall) expr;
			List<AstNode> args = call.getArguments();
			Node target = call.getTarget();
			boolean transformed = false;
			if(!(target instanceof PropertyGet) || !( ContextCheck.isAtomicExpression(((PropertyGet)target).getLeft()) )) {
				if(target instanceof PropertyGet) {
					PropertyGet pg = (PropertyGet) target;
					AstNode left = pg.getLeft();
					Node newTarget = transform(false, left, acc, gs);
					pg.setLeft((AstNode) newTarget);
					transformed = true;
				} else {
					Node newTarget = transform(false, (AstNode) target, acc, gs);
					call.setTarget((AstNode) newTarget);
					//transformed = true;
				}
			}
			List<AstNode> newArgs = new LinkedList<AstNode>();
			for(AstNode arg : args) {
				Node node = transform(false, arg, acc, gs);
				Name name = ASTFactory.name(newName());
				Assignment assign = ASTFactory.assign(ASTFactory.name("var " + name.getString()), (AstNode) node);
				acc.add(ASTFactory.exprStmt(assign));
				newArgs.add(name);
			}
			call.setArguments(newArgs);
			if(transformed) {
				return call;
			} else {
				Name name = ASTFactory.name(newName());
				Assignment push = ASTFactory.assign(ASTFactory.name("var " + name.getString()), expr);
				//Name name = ASTFactory.name("__retTemp.shift()");
				//FunctionCall push = ASTFactory.call("__retTemp.push", expr);
				//Assignment assign = ASTFactory.assign(name, expr);
				acc.add(ASTFactory.exprStmt(push));
				//if(expr instanceof NewExpression) {
				//	acc.add(ASTFactory.exprStmt(ASTFactory.call("_pages.add", name)));
				//}
				return name;
			}
		}
		else if(expr instanceof ArrayLiteral) {
			ArrayLiteral literal = (ArrayLiteral) expr;
			List<AstNode> args = literal.getElements();
			List<AstNode> newArgs = new LinkedList<AstNode>();
			for(AstNode arg : args) {
				Node node = transform(false, arg, acc, gs);
				Name name = ASTFactory.name(newName());
				Assignment assign = ASTFactory.assign(ASTFactory.name("var " + name.getString()), (AstNode) node);
				acc.add(ASTFactory.exprStmt(assign));
				newArgs.add(name);
			}
			literal.setElements(newArgs);
			Name name = ASTFactory.name(newName());
			Assignment assign = ASTFactory.assign(name("var " + name.getString()), expr);
			acc.add(ASTFactory.exprStmt(assign));
			return name;
		} else if(expr instanceof ObjectLiteral) {
			ObjectLiteral literal = (ObjectLiteral) expr;
			List<ObjectProperty> args = literal.getElements();
			for(ObjectProperty arg : args) {
				Node newExpr = transform(false, arg.getRight(), acc, gs);
				Name name = ASTFactory.name(newName());
				Assignment assign = ASTFactory.assign(ASTFactory.name("var " + name.getString()), (AstNode) newExpr);
				acc.add(ASTFactory.exprStmt(assign));
				arg.setRight((AstNode) name);
			}
			Name name = ASTFactory.name(newName());
			Assignment assign = ASTFactory.assign(name("var " + name.getString()), expr);
			acc.add(ASTFactory.exprStmt(assign));
			return name;
		} else if(expr instanceof UnaryExpression) {
			UnaryExpression unary = (UnaryExpression) expr;
			Node newExpr = transform(false, unary.getOperand(), acc, gs);
			unary.setOperand((AstNode) newExpr);
			return unary;
		} else if(expr instanceof FunctionNode) {
			FunctionNode fn = (FunctionNode) expr;
			Block block = (Block) fn.getBody();
			Iterator<Node> children = block.iterator();
			while(children.hasNext()) {
				AstNode child = (AstNode) children.next();
				visit(child);
			}
			return fn;
		}
		else {
			return expr;
		}
	}
	
	@Override
	public boolean visit(AstNode astNode) {
		List<Node> newStmts = new LinkedList<Node>();
		List<Node> gs = new LinkedList<Node>();
		AstNode expr = null;
		if((astNode instanceof ExpressionStatement) || (astNode instanceof ReturnStatement)) {
			if(astNode instanceof ExpressionStatement) {
				ExpressionStatement exprStmt = (ExpressionStatement) astNode;
				expr = exprStmt.getExpression();
			} else {
				ReturnStatement exprStmt = (ReturnStatement) astNode;
				expr = exprStmt.getReturnValue();
			}
			if(ContextCheck.isAtomicExpression(expr)) {
				return true;
			} else if(expr instanceof InfixExpression) {
				if(astNode.getParent() != null) {
					InfixExpression infix = (InfixExpression) expr;
					Node left = transform(ContextCheck.isAssignmentOp(infix.getType()), infix.getLeft(), newStmts, gs);
					Node right = transform(false, infix.getRight(), newStmts, gs);
					infix.setLeft((AstNode) left);
					infix.setRight((AstNode) right);
					tx.insertBefore(astNode.getParent(), astNode, newStmts);
					tx.insertAfter(astNode.getParent(), astNode, gs);
				}
			} else if(expr instanceof VariableDeclaration) {
				VariableDeclaration decl = (VariableDeclaration) expr;
				VariableInitializer init = decl.getVariables().get(0);
				Node right = transform(false, init.getInitializer(), newStmts, gs);
				init.setInitializer((AstNode) right);
				tx.insertBefore(astNode.getParent(), astNode, newStmts);
				tx.insertAfter(astNode.getParent(), astNode, gs);
			} else if(expr instanceof FunctionCall) {
				FunctionCall call = (FunctionCall) expr;
				List<AstNode> args = call.getArguments();
				List<AstNode> newArgs = new LinkedList<AstNode>();
				for(AstNode arg : args) {
					newArgs.add((AstNode) transform(false, arg,newStmts, gs));
				}
				call.setArguments(newArgs);
				Node left = transform(true, call.getTarget(), newStmts, gs);
				call.setTarget((AstNode) left);
				tx.insertBefore(astNode.getParent(), astNode, newStmts);
				tx.insertAfter(astNode.getParent(), astNode, gs);
			}
		} else if(astNode instanceof IfStatement) {
			IfStatement ifStmt = (IfStatement) astNode;
			AstNode condition = ifStmt.getCondition();
			Node newCondition = transform(false, condition, newStmts, gs);
			ifStmt.setCondition((AstNode) newCondition);
			tx.insertBefore(astNode.getParent(), astNode, newStmts);
			tx.insertAfter(astNode.getParent(), astNode, gs);
		}
		return true;
	}
	
}